"""
嵌入模型管理器

这个模块负责管理各种嵌入模型的初始化、配置和调用。
支持多种嵌入模型提供商，包括OpenAI、本地模型等。
"""

import logging
from typing import Optional, Dict, Any, List
from langchain_core.embeddings import Embeddings
from langchain_openai import OpenAIEmbeddings
from langchain_community.embeddings import HuggingFaceEmbeddings
from config.settings import settings

logger = logging.getLogger(__name__)

class EmbeddingManager:
    """
    嵌入模型管理器类
    
    负责管理和提供对各种嵌入模型的统一访问接口。
    支持模型切换、参数调整等功能。
    """
    
    def __init__(self):
        """初始化嵌入模型管理器"""
        self.current_model: Optional[Embeddings] = None
        self.available_models: Dict[str, Dict[str, Any]] = {}
        self._initialize_models()
    
    def _initialize_models(self):
        """初始化可用的嵌入模型"""
        logger.info("正在初始化嵌入模型...")
        
        # 初始化硅基流动嵌入模型
        if settings.OPENAI_API_KEY:
            self.available_models.update({
                "BAAI/bge-m3": {
                    "class": OpenAIEmbeddings,
                    "params": {
                        "model": "BAAI/bge-m3",
                        "api_key": settings.OPENAI_API_KEY,
                        "base_url": settings.OPENAI_API_BASE,
                        "dimensions": 1024
                    },
                    "description": "BGE-M3 - 多功能、多语言、多粒度的文本嵌入模型，支持三种常见的检索功能"
                },
                "netease-youdao/bce-embedding-base_v1": {
                    "class": OpenAIEmbeddings,
                    "params": {
                        "model": "netease-youdao/bce-embedding-base_v1",
                        "api_key": settings.OPENAI_API_KEY,
                        "base_url": settings.OPENAI_API_BASE,
                        "dimensions": 768
                    },
                    "description": "BCE-embedding-base_v1 - 网易有道开发的双语和跨语言嵌入模型"
                },
                "BAAI/bge-large-zh-v1.5": {
                    "class": OpenAIEmbeddings,
                    "params": {
                        "model": "BAAI/bge-large-zh-v1.5",
                        "api_key": settings.OPENAI_API_KEY,
                        "base_url": settings.OPENAI_API_BASE,
                        "dimensions": 1024
                    },
                    "description": "BGE-large-zh-v1.5 - 智源研究院大型中文文本嵌入模型"
                },
                "BAAI/bge-large-en-v1.5": {
                    "class": OpenAIEmbeddings,
                    "params": {
                        "model": "BAAI/bge-large-en-v1.5",
                        "api_key": settings.OPENAI_API_KEY,
                        "base_url": settings.OPENAI_API_BASE,
                        "dimensions": 1024
                    },
                    "description": "BGE-large-en-v1.5 - 智源研究院大型英文文本嵌入模型"
                },
            })
            logger.info(f"已加载 {len(self.available_models)} 个硅基流动嵌入模型")
        
        # 本地嵌入模型（可选）
        local_models = {
            "bge-small-zh-v1.5": {
                "class": HuggingFaceEmbeddings,
                "params": {
                    "model_name": "BAAI/bge-small-zh-v1.5",
                    "model_kwargs": {"device": "cpu"},
                    "encode_kwargs": {"normalize_embeddings": True}
                },
                "description": "BGE-small-zh-v1.5 - 智源研究院中文嵌入模型(本地版)"
            },
            "bge-base-zh-v1.5": {
                "class": HuggingFaceEmbeddings,
                "params": {
                    "model_name": "BAAI/bge-base-zh-v1.5",
                    "model_kwargs": {"device": "cpu"},
                    "encode_kwargs": {"normalize_embeddings": True}
                },
                "description": "BGE-base-zh-v1.5 - 智源研究院中文嵌入模型(本地基础版)"
            }
        }
        
        # 检查是否可以使用本地模型
        try:
            import torch
            self.available_models.update(local_models)
            logger.info(f"已加载 {len(local_models)} 个本地嵌入模型")
        except ImportError:
            logger.info("未安装torch，跳过本地嵌入模型")
        
        # 设置默认模型 - 优先使用BGE-M3
        if self.available_models:
            default_model_name = getattr(settings, 'DEFAULT_EMBEDDING_MODEL', 'BAAI/bge-m3')
            if default_model_name in self.available_models:
                self.set_model(default_model_name)
            else:
                # 使用BGE-M3作为默认模型
                if 'BAAI/bge-m3' in self.available_models:
                    self.set_model('BAAI/bge-m3')
                    logger.info(f"默认嵌入模型 {default_model_name} 不可用，使用 BAAI/bge-m3")
                else:
                    # 使用第一个可用模型作为默认模型
                    first_model = list(self.available_models.keys())[0]
                    self.set_model(first_model)
                    logger.info(f"BGE-M3不可用，使用 {first_model}")
        else:
            logger.warning("没有可用的嵌入模型")
    
    def set_model(self, model_name: str, **kwargs) -> bool:
        """
        设置当前使用的嵌入模型
        
        Args:
            model_name (str): 模型名称
            **kwargs: 额外的模型参数
            
        Returns:
            bool: 是否设置成功
        """
        if model_name not in self.available_models:
            logger.error(f"嵌入模型 {model_name} 不可用")
            return False
        
        try:
            model_config = self.available_models[model_name].copy()
            
            # 合并额外参数
            if kwargs:
                model_config["params"].update(kwargs)
            
            # 创建模型实例
            model_class = model_config["class"]
            self.current_model = model_class(**model_config["params"])
            
            logger.info(f"已切换到嵌入模型: {model_name}")
            return True
            
        except Exception as e:
            logger.error(f"设置嵌入模型 {model_name} 失败: {str(e)}")
            return False
    
    def get_current_model(self) -> Optional[Embeddings]:
        """
        获取当前嵌入模型实例
        
        Returns:
            Optional[Embeddings]: 当前嵌入模型实例
        """
        return self.current_model
    
    def list_available_models(self) -> Dict[str, str]:
        """
        列出所有可用的嵌入模型
        
        Returns:
            Dict[str, str]: 模型名称和描述的字典
        """
        return {
            name: config["description"] 
            for name, config in self.available_models.items()
        }
    
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """
        嵌入文档列表
        
        Args:
            texts (List[str]): 文本列表
            
        Returns:
            List[List[float]]: 嵌入向量列表
        """
        if not self.current_model:
            raise ValueError("没有设置当前嵌入模型")
        
        try:
            return self.current_model.embed_documents(texts)
        except Exception as e:
            logger.error(f"文档嵌入失败: {str(e)}")
            raise
    
    def embed_query(self, text: str) -> List[float]:
        """
        嵌入查询文本
        
        Args:
            text (str): 查询文本
            
        Returns:
            List[float]: 嵌入向量
        """
        if not self.current_model:
            raise ValueError("没有设置当前嵌入模型")
        
        try:
            return self.current_model.embed_query(text)
        except Exception as e:
            logger.error(f"查询嵌入失败: {str(e)}")
            raise
    
    def get_model_info(self) -> Dict[str, Any]:
        """
        获取当前嵌入模型的信息
        
        Returns:
            Dict[str, Any]: 模型信息
        """
        if not self.current_model:
            return {
                "is_ready": False,
                "error": "没有设置当前嵌入模型"
            }
        
        # 查找当前模型的配置
        for name, config in self.available_models.items():
            if isinstance(self.current_model, config["class"]):
                return {
                    "is_ready": True,
                    "model_name": name,
                    "description": config["description"],
                    "class": config["class"].__name__,
                    "params": config["params"]
                }
        
        return {
            "is_ready": True,
            "error": "无法获取模型信息"
        }
    
    def is_ready(self) -> bool:
        """
        检查嵌入模型是否就绪
        
        Returns:
            bool: 是否就绪
        """
        return self.current_model is not None
    
    def test_embedding(self, test_text: str = "测试文本") -> bool:
        """
        测试嵌入模型是否正常工作
        
        Args:
            test_text (str): 测试文本
            
        Returns:
            bool: 是否正常工作
        """
        if not self.current_model:
            return False
        
        try:
            # 测试嵌入
            embedding = self.embed_query(test_text)
            return len(embedding) > 0
        except Exception as e:
            logger.error(f"嵌入模型测试失败: {str(e)}")
            return False

# 创建全局嵌入模型管理器实例
embedding_manager = EmbeddingManager() 